package com.szjlc.blog.common.utils;

import java.util.Map;
import java.util.concurrent.Callable;
import org.slf4j.MDC;

/**
 * 多线程MDC工具类
 * @author steve
 * @since 2022-11-28 22:12
 */

public class ThreadMDCUtil {

    /** 若有traceId则设置追踪id */
    public static void setTraceIfAbsent() {
        if (MDC.get(TraceIdUtil.TRACE_ID) == null) {
            MDC.put(TraceIdUtil.TRACE_ID, TraceIdUtil.generateTraceId());
        }
    }

    public static <T> Callable<T> wrap(final Callable<T> callable, Map<String, String> context) {
        return () -> {
            // 1.判断当前线程对应MDC的Map是否存在，如果存在则设置Map
            if (context == null) {
                MDC.clear();
            }
            else {
                MDC.setContextMap(context);
            }
            // 2.设置MDC中的traceId值，不存在则新生成，如果是子线程，MDC中traceId不为null
            setTraceIfAbsent();

            try {
                // 3.执行线程方法
                return callable.call();
            }
            finally {
                // 最后无论线程如何了，都要清空MDC防止内存泄漏(占用),MDC底层ThreadLocal
                MDC.clear();
            }
        };
    }

    /* 代码说明：
        1.判断当前线程对应MDC的Map是否存在，如果存在则设置Map
        2.设置MDC中的traceId值，不存在则新生成，如果是子线程，MDC中traceId不为null
        3.执行run方法 */
    public static Runnable wrap(Runnable runnable, Map<String, String> context) {
        return () -> {
            if (context == null) {
                MDC.clear();
            }
            else {
                MDC.setContextMap(context);
            }

            setTraceIfAbsent();

            try {
                runnable.run();
            }
            finally {
                MDC.clear();
            }
        };
    }
}
